package cn.swing.main.srv.cv.onnx;

import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.core.MatOfRect2d;
import org.opencv.core.Rect2d;
import org.opencv.core.Scalar;
import org.opencv.core.Size;
import org.opencv.dnn.Dnn;
import org.opencv.dnn.Net;
import org.opencv.videoio.VideoCapture;

import org.opencv.core.Core;
import org.opencv.core.MatOfFloat;
import org.opencv.core.MatOfInt;
import org.opencv.core.Point;
import org.opencv.imgproc.Imgproc;
import org.opencv.videoio.VideoWriter;
import org.opencv.videoio.Videoio;
import org.opencv.highgui.HighGui;

import java.util.ArrayList;
import java.util.List;

public class OnnxModelTrace {

    static {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);
    }

    public void run() {
        // 视频路径
        String path = "D:\\swing\\3f5bfd1d-4d33-4c3d-846a-24dd48855f4c-1.mp4";
        VideoCapture capture = new VideoCapture(path);
        if (!capture.isOpened()) {
            System.out.println("Error: Could not open video capture.");
            return;
        }

        // 模型路径
        String modelPath = "D:\\swing\\yolo11n.onnx";
        Net net = Dnn.readNetFromONNX(modelPath);

        // 输出视频设置
        VideoWriter videoWriter = null;
        int frameWidth = (int) capture.get(Videoio.CAP_PROP_FRAME_WIDTH);
        int frameHeight = (int) capture.get(Videoio.CAP_PROP_FRAME_HEIGHT);
        double fps = capture.get(Videoio.CAP_PROP_FPS);
        videoWriter = new VideoWriter("D:\\swing\\output_video.mp4",
                VideoWriter.fourcc('m', 'p', '4', 'v'), fps,
                new Size(frameWidth, frameHeight));

        Mat frame = new Mat();
        Mat blobFrame = new Mat();
        MatOfFloat confidences = new MatOfFloat();
        MatOfInt classIds = new MatOfInt();
        List<Rect2d> boxes = new ArrayList<>();

        while (capture.read(frame)) {
            if (frame.empty()) {
                break;
            }

            // 创建 Blob
            blobFrame = Dnn.blobFromImage(frame, 1.0 / 255.0, new Size(416, 416), new Scalar(0, 0, 0), true, false,
                    CvType.CV_32F);

            // 设置输入
            net.setInput(blobFrame);

            // 前向传播
            List<Mat> outs = new ArrayList<>();
            net.forward(outs, getOutputsNames(net));

            // 解析输出
            postprocess(frame, outs, confidences, classIds, boxes);

            // 绘制检测结果
            for (int i = 0; i < boxes.size(); i++) {
                Rect2d box = boxes.get(i);
                int classId = (int) classIds.get(i, 0)[0];
                double confidence = confidences.get(i, 0)[0];
                String label = "Class " + classId + ": " + String.format("%.2f", confidence);
                Imgproc.rectangle(frame, box.tl(), box.br(), new Scalar(0, 255, 0), 2);
                Imgproc.putText(frame, label, new Point(box.tl().x, box.tl().y - 5), Imgproc.FONT_HERSHEY_SIMPLEX,
                        0.5, new Scalar(0, 255, 0), 2);
            }

            // 显示结果帧
            HighGui.imshow("YOLO ONNX Detection", frame);

            // 写入输出视频
            videoWriter.write(frame);

            // 按下 'q' 键退出
            if (HighGui.waitKey(30) == 'q') {
                break;
            }
        }

        capture.release();
        videoWriter.release();
        HighGui.destroyAllWindows();
    }

    private void postprocess(Mat frame, List<Mat> outs, MatOfFloat confidences, MatOfInt classIds, List<Rect2d> boxes) {
        float confThreshold = 0.5f;
        float nmsThreshold = 0.4f;
        int frameHeight = frame.height();
        int frameWidth = frame.width();

        for (Mat out : outs) {
            float[] data = new float[(int) (out.total() * out.channels())];
            out.get(0, 0, data);
            for (int i = 0; i < data.length; i += 85) {
                float confidence = data[i + 4];
                if (confidence > confThreshold) {
                    float[] box = new float[4];
                    for (int j = 0; j < 4; j++) {
                        box[j] = data[i + 5 + j];
                    }
                    float centerX = box[0] * frameWidth;
                    float centerY = box[1] * frameHeight;
                    float width = box[2] * frameWidth;
                    float height = box[3] * frameHeight;
                    float left = centerX - width / 2;
                    float top = centerY - height / 2;

                    classIds.push_back(new MatOfInt(getMaxClassId(data, i + 5)));
                    confidences.push_back(new MatOfFloat(confidence));
                    boxes.add(new Rect2d((int) left, (int) top, (int) width, (int) height));
                }
            }
        }

        MatOfRect2d matOfRect = new MatOfRect2d();
        matOfRect.fromList(boxes);

        // 非极大值抑制
        MatOfInt indices = new MatOfInt();
        Dnn.NMSBoxes(matOfRect, confidences, confThreshold, nmsThreshold, indices);

        // 提取最终的检测结果
        for (int i = 0; i < indices.size().height; i++) {
            int idx = (int) indices.get(i, 0)[0];
            classIds.put(i, 0, classIds.get(idx, 0));
            confidences.put(i, 0, confidences.get(idx, 0));
            boxes.set(i, boxes.get(idx));
        }
    }

    private int getMaxClassId(float[] data, int start) {
        int maxClassId = 0;
        float maxConfidence = 0;
        for (int i = 0; i < 80; i++) {
            float confidence = data[start + i];
            if (confidence > maxConfidence) {
                maxConfidence = confidence;
                maxClassId = i;
            }
        }
        return maxClassId;
    }

    private List<String> getOutputsNames(Net net) {
        List<String> names = new ArrayList<>();
        List<Integer> outLayers = net.getUnconnectedOutLayers().toList();
        List<String> layersNames = net.getLayerNames();
        for (int i = 0; i < outLayers.size(); i++) {
            names.add(layersNames.get(outLayers.get(i) - 1));
        }
        return names;
    }

    public static void main(String[] args) {
        OnnxModelTrace detector = new OnnxModelTrace();
        detector.run();
    }

}

